Populating next right pointers in each node II

Time: O(N); Space: O(1); hard

Follow up for problem “Populating Next Right Pointers in Each Node”.

What if the given tree could be any binary tree? Would your previous solution still work?

Note:

• You may only use constant extra space.

Example 1:

Input:

    1
   /  \
  2    3
 / \    \
4   5    7

Output:

    1 -> NULL
   /  \
  2 -> 3 -> NULL
 / \    \
4-> 5 -> 7 -> NULL

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
        self.next = None

    def __repr__(self):
        if self is None:
            return "Nil"
        else:
            return "{} -> {}".format(self.val, repr(self.next))

class Solution1(object):
    def connect(self, root):
        """
        :type root: TreeNode
        :rtype: None
        """
        head = root
        while head:
            prev, cur, next_head = None, head, None
            while cur:
                if next_head is None:
                    if cur.left:
                        next_head = cur.left
                    elif cur.right:
                        next_head = cur.right

                if cur.left:
                    if prev:
                        prev.next = cur.left
                    prev = cur.left

                if cur.right:
                    if prev:
                        prev.next = cur.right
                    prev = cur.right

                cur = cur.next
            head = next_head

if __name__ == "__main__":
    root, root.left, root.right = TreeNode(1), TreeNode(2), TreeNode(3)
    root.left.left, root.left.right, root.right.right = TreeNode(4), TreeNode(5), TreeNode(7)
    Solution1().connect(root)
    print(root)
    print(root.left)
    print(root.left.left)

1 -> None
2 -> 3 -> None
4 -> 5 -> 7 -> None

See also:

https://leetcode.com/problems/populating-next-right-pointers-in-each-node-ii